#!/usr/bin/env python
"""
usage:
    %prog <pkdb filename> [options]

A utility for managing public key database (PKDB). a PKDB contains a private
key, followed by a list of "authorized" public keys, i.e., keys which we trust.
this is basically a unification of SSH's id_rsa and authorized_keys into a 
single file
"""
import sys
import os
from optparse import OptionParser
from rpyc.utils.security import RSAKey, PubkeyDB


parser = OptionParser(usage=__doc__)
parser.add_option("-c", "--create", dest="create", action="store_true", 
    default = False, help="creates a new pkdb file with a new private key")
parser.add_option("-p", "--public", dest="public", action="store_true", 
    default = False, help="returns the public portion of the pkdb's private key")
parser.add_option("-a", "--add", dest="add", action="store", default = None,
    metavar = "KEY", help="adds the given KEY (public key) to the pkdb")
parser.add_option("-b", "--label", dest="label", action="store", default = None,
    metavar = "LABEL", help="sets the label for -a, -f, -u, -r; if not given, "
    "you will be prompted for a label")
parser.add_option("-r", "--rename", dest="rename_label", action="store", default = None,
    metavar = "LABEL", help="renames LABEL to a new name")
parser.add_option("-u", "--dup", dest="dup_label", action="store", default = None,
    metavar = "LABEL", help="duplicates the key of LABEL with a new label")
parser.add_option("-f", "--add-file", dest="add_file", action="store", default = None,
    metavar="FILENAME", help="same as -a, but reads the key from FILENAME")
parser.add_option("-d", "--delete-label", dest="delete_label", action="store",
    metavar = "LABEL", default = None, help="deletes the key associated with LABEL")
parser.add_option("-k", "--delete-key", dest="delete_key", action="store", 
    metavar = "KEY", default = None, help="deletes the key KEY from the pkdb")
parser.add_option("-l", "--list", dest="list", action="store_true",
    help="lists all public keys contained in the pkdb file")


if __name__ == "__main__":
    options, args = parser.parse_args()
    if len(args) != 1:
        parser.error("missing key db file")
    
    actions = [options.create, options.public, options.add, options.add_file, 
        options.delete_label, options.delete_key, options.list, options.dup_label,
        options.rename_label]
    if sum(bool(opt) for opt in actions) > 1:
        parser.error("options -a, -d, -k, -f, -l, -r, -u are mutually exclusive")
    
    filename = args[0]
    
    if options.create:
        db = PubkeyDB.create(filename)
        db.save()
        print "initialized the pkdb with a new private key. The public key is:"
        print db.private_key.dump_public()
        sys.exit()

    if os.path.exists(filename):
        db = PubkeyDB.load(filename)
    else:
        db = PubkeyDB(None, {}, filename)
    
    if options.public:
        print db.private_key.dump_public()
    elif options.list:
        if not db:
            print "pkdb file is empty"
        else:
            print "contained public keys:"
            for label, pk in db:
                print " * %-20s: %s" % (label, pk.dump_public())
    elif options.add:
        rawkey = options.add
        pk = RSAKey.load(rawkey)
        if not options.label:
            options.label = raw_input("Enter a label for this key: ").strip()
        db.add_public_key(options.label, pk)
        db.save()
        print "added key to pkdb file"
    elif options.add_file:
        rawkey = open(options.add_file, "r").read()
        pk = RSAKey.load(rawkey)
        if not options.label:
            options.label = raw_input("Enter a label for this key: ").strip()
        db.add_public_key(options.label, pk)
        db.save()
        print "added key to pkdb file"
    elif options.delete_key:
        rawkey = options.delete_key
        pk = RSAKey.load(rawkey)
        db.remove_public_key_by_key(pk)
        db.save()
        print "removed key from pkdb file"
    elif options.delete_label:
        label = options.delete_label
        db.remove_public_key(label)
        db.save()
        print "removed key from pkdb file"
    elif options.rename_label:
        old_label = options.rename_label
        if not options.label:
            options.label = raw_input("Enter the new label: ").strip()
        pk = db.get(old_label)
        db.remove_public_key(old_label)
        db.add_public_key(options.label, pk)
        db.save()
        print "renamed %s to %s" % (old_label, options.label)
    elif options.dup_label:
        orig_label = options.dup_label
        if not options.label:
            options.label = raw_input("Enter the new label: ").strip()
        pk = db.get(orig_label)
        db.add_public_key(options.label, pk)
        db.save()
        print "duplicated %s as %s" % (orig_label, options.label)
    else:
        parser.error("no action specified")



